#-*- coding:utf-8 -*-
"""
Created on Tue Oct 24 08:25:40 2017
显示标准pascal voc 标注数据
@author: aia_zxl
"""
import argparse
import shutil
import os
import sys
import xml.etree.ElementTree as ET
import cv2


def parse_pascal_xml(xml_filename, filter_diff=False):
    """
    parse object information from xml file. return image name.
    :param xml_filename: xml file
    :param filter_diff: wether or not filter out difficult objects
    :return: image name and object information.
    """
    tree = ET.parse(xml_filename)
    root = tree.getroot()
    obj_info = []
    img_name = root.find('filename').text
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        object_class = obj.find('name').text
        if filter_diff is True and int(difficult) == 1:
            continue
        xmlbox = obj.find('bndbox')
        rect = (float(xmlbox.find('xmin').text),
                float(xmlbox.find('ymin').text),
                float(xmlbox.find('xmax').text),
                float(xmlbox.find('ymax').text))
        obj_info.append({'object_class': object_class, 'rect': rect, 'difficult': int(difficult)})
    return img_name, obj_info


def show_data(root_path, list_file, label_strs, show_img=False):
    """
    save images with object marked.
    :param root_path: voc root path
    :param list_file: text file contain all file lists
    :param label_strs: object class name
    :param show_img: show image if true
    :return: None
    """
    wait_sec = 0
    # cv2.namedWindow('Image', cv2.WINDOW_FREERATIO)
    xml_num = len(os.listdir(os.path.join(root_path, "Annotations")))
    save_path = os.path.join(root_path, 'img_with_rectangle')
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    else:
        shutil.rmtree(save_path)
        os.makedirs(save_path)

    file_names = open(list_file, 'r').readlines()
    processed_num = 0
    for file_name in file_names:
        xml_file = os.path.join(root_path, 'Annotations/{}.xml'.format(file_name.strip()))
        img_name, obj_info = parse_pascal_xml(xml_file)
        # file_nm = a_label.strip()
        img_file_full_path = os.path.join(root_path, 'JPEGImages', img_name)
        img = cv2.imread(img_file_full_path)

        for _, data in enumerate(obj_info):
            class_id = -1
            if data['object_class'] in label_strs:
                class_id = label_strs.index(data['object_class'])
            else:
                print("unknow class name: {}".format(data["object_class"]))
                continue

            if class_id == 0:
                color = (0, 0, 255)
            else:
                color = (0, 255, 0)

            rect = data['rect']
            cv2.rectangle(img, (int(rect[0]), int(rect[1])),
                          (int(rect[2]), int(rect[3])), color, 2)
            cv2.putText(img, data['object_class'],
                        (int(rect[0]), int(rect[1]) - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
            #显示文件名
            # cv2.putText( img,file_nm,(15,15), cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0) )

        processed_num += 1
        print("Processed: {}, Total: {}".format(processed_num, xml_num))
        cv2.imwrite(os.path.join(save_path, img_name), img)
        if show_img:
            cv2.imshow('Image', img)
            key = cv2.waitKey(wait_sec)
            if key == 32:
                wait_sec = 1-wait_sec
            elif key == ord('q') or key == ord('Q'):
                # cv2.destroyAllWindows()
                return

    # cv2.destroyAllWindows()


if __name__ == '__main__':

    # LABELS = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
    #           "car", "cat", "chair", "cow", "diningtable", "dog", "horse",
    #           "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
    LABELS = ["tank", "plane", "ship"]
    # colors = plt.cm.hsv(np.linspace(0, 1, 1+len(labels))).tolist()
    try:
        ROOT_DIR = sys.argv[1]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/zongshiban_data/3/ccj/dataset/plane_ship_tank_VOC",
                            help="VOC dataset root directory")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
    FILE_LIST = os.path.join(ROOT_DIR, "ImageSets/Main/train.txt")
    show_data(ROOT_DIR, FILE_LIST, LABELS, False)
